#!/usr/bin/env python3
"""
TAPD 季度总结自动生成系统 - 主入口文件
基于 AutoGen 框架，集成 DeepSeek 大模型和 MCP 工具
"""

import asyncio
import os
import sys
from datetime import datetime, timedelta
from typing import Dict, List, Optional

# 添加项目根目录到 Python 路径
sys.path.append(os.path.dirname(os.path.dirname(__file__)))

from dotenv import load_dotenv
from autogen_agentchat.agents import AssistantAgent
from autogen_agentchat.teams import RoundRobinGroupChat

from src.agents.data_scraper_agent import DataScraperAgent
from src.agents.expert_agent import ExpertAgent
from src.agents.assistant_agent import AssistantReportAgent
from src.config.settings import load_config
from src.utils.model_client import create_deepseek_client


def get_quarter_dates(year: int, quarter: int) -> tuple:
    """获取指定年份季度的开始和结束日期"""
    if quarter == 1:
        start_date = datetime(year, 1, 1)
        end_date = datetime(year, 3, 31)
    elif quarter == 2:
        start_date = datetime(year, 4, 1)
        end_date = datetime(year, 6, 30)
    elif quarter == 3:
        start_date = datetime(year, 7, 1)
        end_date = datetime(year, 9, 30)
    elif quarter == 4:
        start_date = datetime(year, 10, 1)
        end_date = datetime(year, 12, 31)
    else:
        raise ValueError("季度必须是1-4之间的数字")

    return start_date, end_date


async def main():
    """主函数"""
    # 加载环境变量
    load_dotenv()

    # 加载配置
    config = load_config()

    print("🚀 TAPD 季度总结自动生成系统启动")
    print("=" * 50)

    # 获取用户输入
    try:
        year = int(input("请输入年份 (例如: 2024): ").strip() or str(datetime.now().year))
        quarter = int(input("请输入季度 (1-4): ").strip())

        if quarter not in [1, 2, 3, 4]:
            raise ValueError("季度必须是1-4之间的数字")

        start_date, end_date = get_quarter_dates(year, quarter)
        print(f"\n📅 分析时间范围: {start_date.strftime('%Y-%m-%d')} 到 {end_date.strftime('%Y-%m-%d')}")

    except ValueError as e:
        print(f"❌ 输入错误: {e}")
        return

    try:
        # 初始化模型客户端
        model_client = create_deepseek_client(
            api_key=config.deepseek_api_key,
            base_url=config.deepseek_base_url
        )

        # 初始化智能体
        print("\n🤖 初始化智能体...")

        # 数据抓取智能体
        scraper_agent = DataScraperAgent(
            name="数据抓取员",
            model_client=model_client,
            config=config
        )

        # 专家分析智能体
        expert_agent = ExpertAgent(
            name="数据分析专家",
            model_client=model_client,
            config=config
        )

        # 助手报告智能体
        assistant_agent = AssistantReportAgent(
            name="报告生成助手",
            model_client=model_client,
            config=config
        )

        # 创建团队协作
        team = RoundRobinGroupChat([scraper_agent, expert_agent, assistant_agent])

        # 构建任务描述
        task_description = f"""
        请完成 {year} 年第 {quarter} 季度的 TAPD 工作总结报告生成任务：

        1. 数据抓取员：请从 TAPD 系统抓取 {start_date.strftime('%Y-%m-%d')} 到 {end_date.strftime('%Y-%m-%d')} 期间的工作项数据
        2. 数据分析专家：对抓取的数据进行深度分析，提取关键指标和趋势
        3. 报告生成助手：基于分析结果生成专业的季度总结报告

        请按顺序执行，确保每个步骤的输出为下一步提供必要的信息。
        """

        print("\n📊 开始执行季度总结任务...")
        print("-" * 30)

        # 执行任务
        result = await team.run(task=task_description)

        print("\n✅ 季度总结报告生成完成！")
        print("=" * 50)

        # 保存结果
        output_dir = f"output/{year}_Q{quarter}"
        os.makedirs(output_dir, exist_ok=True)

        # 保存完整对话记录
        with open(f"{output_dir}/conversation_log.md", "w", encoding="utf-8") as f:
            f.write(f"# {year} 年第 {quarter} 季度 TAPD 总结报告生成记录\n\n")
            f.write(f"生成时间: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n\n")
            f.write("## 完整对话记录\n\n")
            for message in result.messages:
                f.write(f"**{message.source}**: {message.content}\n\n")

        print(f"📁 完整记录已保存到: {output_dir}/conversation_log.md")

    except Exception as e:
        print(f"❌ 执行过程中出现错误: {str(e)}")
        import traceback
        traceback.print_exc()


if __name__ == "__main__":
    asyncio.run(main())